#ifndef __AK_SESSION__
#define __AK_SESSION__


namespace xm
{
	class DetectSession : public IKDetectSession, UnkBase
	{
	protected:
		struct MyDetectResult
		{
			DETLEV			Level;
			UINT32			SigId;
			std::string		FullName;

			MyDetectResult()
			{
				Reset();
			}

			~MyDetectResult()
			{
			}
			VOID	Reset()
			{
				FullName.clear();
				Level = DLEV_NOT_MALW;
				SigId = 0;
			}
			HRESULT SetResult( DetectResult& result )
			{
				if( result.Level <= Level ) return E_UNEXPECTED;
				Level = result.Level;
				SigId = result.SigId;
				FullName = result.Name;
				return S_OK;
			}
		};
	protected:
		iptr<IKXoTable>			GlobalTable;
		iptr<IKXoTable>			SessionTable;
		iptr<IKFileStream>		RawFile;
		iptr<IKBStream>			File;
		iptr<IKArchive>			Archive;
		MyDetectResult			Result;
		UINT32					Options;
		HRESULT					TreatResult;
		typedef std::vector<MyDetectResult>	MyResults_t;
		MyResults_t				Results;
	public:
		UNKNOWN_IMP1(IKDetectSession);
		DetectSession() : Options(0), TreatResult(S_OK) {};
		HRESULT init_class( IUnknown * got, IUnknown * outer )
		{
			if( outer ) return E_FAIL;
			RFAILED( SessionTable.instantiation( PROGID_XoTable ) );
			GlobalTable = got;
			return S_OK;
		}
		STDMETHOD(Open)( IUnknown * pstm, UINT32 uOptions )
		{
			Close();
			RawFile = pstm;
			if( !RawFile ) return E_NOINTERFACE;
			File = RawFile;
			Options = uOptions;
			return S_OK;
		}
		STDMETHOD(SelectFormat)( FMTID fmtid, UINT32 uOpenFlags = -1 )
		{
			Archive.dispose();
			if( !RawFile ) return E_UNEXPECTED;

			iptr<IKArchive>	arch( (FileType)fmtid );
			if( !arch ) return E_FAIL;

			optr<IKBStream> astm;
			if( SUCCEEDED( XmCreateAttachStream( fmtid, RawFile, astm.pp() ) ) )
				File = astm;
			
			RFAILED( arch->Open( File, fmtid, uOpenFlags ) );

			Archive = arch;
			
			return S_OK;
		}
		STDMETHOD(Close)()
		{
			File.dispose();
			Archive.dispose();
			Result.Reset();
			Options = 0;
			TreatResult = S_OK;
			if( SessionTable ) SessionTable->Clear();
			return S_OK;
		}
		STDMETHOD_(IKXoTable*,GlobalXoTable)()
		{
			return GlobalTable;
		}
		STDMETHOD_(IKXoTable*,SessionXoTable)()
		{
			return SessionTable;
		}
		STDMETHOD_(UINT32,DetectOptions)()
		{
			return Options;
		}
		STDMETHOD_(IKFileStream*,RawStream)()
		{
			return RawFile;
		}
		STDMETHOD_(IKBStream*,TargetStream)()
		{
			return File;
		}
		STDMETHOD_(IKArchive*,CurrentArchive)()
		{
			return Archive;
		}
		STDMETHOD(PushResult)( DetectResult & result, IKDetector * det)
		{
			if( !det || !result.SigId || !result.Name ) 
				return E_INVALIDARG;
			if( Options & DOSCAN_LOGDR )
			{
				MyDetectResult dr; 
				dr.SetResult( result );
				Results.push_back( dr );
			}
			return Result.SetResult( result );
		}
		STDMETHOD(GetResult)( DetectResult& result, IKDetector** pp )
		{
			if( Result.Level < DLEV_MALW_HEUR ) 
				return E_UNEXPECTED;
			result.Level = Result.Level;
			result.Name = Result.FullName.c_str();
			result.SigId = Result.SigId;
			if( !pp ) return S_OK;
			return S_OK;
		}
		STDMETHOD(SetTreatResult)( HRESULT hrTreat )
		{
			TreatResult = hrTreat;
			return S_OK;
		}
		STDMETHOD_(SIZE_T,EnumResult)(DetectResult * lpResults, SIZE_T cWant )
		{
			size_t cnt = Results.size();
			if( !cWant ) return cnt;
			if( !lpResults ) return 0;
			cnt = MIN(cWant,cnt);
			for( size_t i = 0; i < cnt; ++ i )
			{
				lpResults[i].Level = Results[i].Level;
				lpResults[i].SigId = Results[i].SigId;
				lpResults[i].Name = Results[i].FullName.c_str();
			}
			return cnt;
		}
		STDMETHOD( SetAttachedStream )( IKBStream * bso )
		{
			File = bso;
			return S_OK;
		}
	};
};



#endif